%% Script for synthesizing CGXHs for complex-amplitude control and simulation of FTH experiments 
% Double-Constraint Gerchberg-Saxton CGH Synthesis;
% 
% Written by Kahraman Keskinbora kahraman@mit.edu keskinbora@is.mpg.de
% between % 2018 - 2021 for the paper "Maskless Fourier Transform Holography"
%
% References:
% Chang et al., Appl. Opt. 54, 6994-7001 (2015)
% Voelz, Computational fourier optics: a MATLAB tutorial (SPIE, 2011)
%
%
%% Initialize the workspace

close all;
clear all;
clc;

load cmap.mat % this is a color scheme used in some of the plots

%% Units;

meters      =1;
microns     =1e-6 * meters;
nanometers  =1e-9 * meters;
cm          =10^-2 * meters;

seconds     =1;
miliseconds =1e-6 * seconds;
nanoseconds =1e-9 * seconds;

hertz       =1/seconds;
kilohertz   =1e6 * hertz;
gigahertz   =1e9 * hertz;
terahertz   =1e12* hertz;
petahertz   =1e15* hertz;

eV          =1;
keV         =eV * 1000;
meV         =eV / 1000;


%% Constants;

c0          = 299792458 * meters/seconds;
e0          = 8.854187817e-12 * 1/meters;
u0          = 1.2566370614e-6 * 1/meters;

%% Sample Definition
% X30-30-2_SS_20um8bit_1kpixRescaled.bmp
% Magnetic_stripe_domains_binary2.bmp

samplefile  = 'X30-30-2_SS_20um8bit_1kpixRescaled.bmp';
% Please correct this path for your computer
path        = 'D:\PostDoc\PostDoc_Work\Data&Analysis\Gerchberg-Saxton CGH\' 
test        = double(imread(strcat(path,samplefile)));

test        = test./max(max(test));

%% Beamstop & OSA Yes/No

beam_stop   = 2;                            % BS & aperture 1 for Yes, 2 for No
Use_OSA     = 1;                            % 1 for Yes, 2 for No

%% Target Definition 
% Target is the desired structured illumination pattern at the focal plane
     
L           = 200*microns;          % Initial Sim bounding box
Nx          = 4000;                 % Number of sampling points in x and y
dx          = L/Nx;                 % Step size
x           = -L/2:dx:L/2-dx;       % Vector x
y           = x;
[X,Y]       = meshgrid(x,y);        % Grid
En          = 1200;                 % Design photon energy
lambda      = (1239.840/En)*nanometers; % Wavelenth
f           = 1.2*cm;               % Design focal length
seperX      = -2*microns;           % Seperation between the ref and sample beam
seperY      = -2*microns;

% Rectangular Sample Beam
fov_FW      = 2*microns;                    % Sample beam FOV width               
fov_hw      = fov_FW/2;                     % FOV half-width
Fov_spot    = rect((X+seperX)./(2*fov_hw)).*rect((Y-seperY)./(2*fov_hw));

% Reference Beam
Ref_spot_D  = dx;
Ref_spot    = circle(Ref_spot_D,0,0,x);

IR          = 150;                          % Intensity ratio of reference spot to sample beam
target_abs  = Fov_spot+IR.*Ref_spot;
target_phase= exp(1i.*pi*(Fov_spot+Ref_spot));
target      = target_abs.*target_phase;     % Target pattern

CGXH_w      = 100*microns;                      % CGXH optic size
support     = rect(X./CGXH_w).*rect(Y./CGXH_w);     % CGXH support



%% Double-Constraint Gerchberg-Saxton Algorithm

G=zeros(Nx);                                % Initate the Input Field;

PWI     = 1;                                  % Plane wave illumination intensity

iterations = 20;                              % Number of iterations for DCGH
source=G;
counter=1;
tc=0;

% DC-GSA parameters

alpha = 1;
beta  = 0.01;
gamma = 0.1;
Npix   = Nx.^2;  
z      = f;

 for i=1:1:iterations           
    tic;
    A=support.*exp(1i*angle(source));
    B=propRS2(A,L,lambda,z);
    C1=(alpha.*abs(target)-beta.*abs(B)).*exp(1i.*angle(target));
    C2=(gamma.*abs(B)).*exp(1i.*angle(B));
    C=C1+C2;
    D=propRS2(C,L,lambda,-z);
    source=D;
    t(i)=toc;
    tc=tc+t(i);

    RMS(i)=error_function_Miao(abs(B),abs(target))
    figure(1)
    plot((RMS))
    drawnow

 end

RMS(i)

% Binarization of the synthesized hologram

holo        = (angle(A)+pi);
holo2       = angle(A);
holo        = holo./max(max(holo));
DC          = 0.5;
binary_CGH  = (0).*(holo<DC)+1.*(holo>=DC);


%% Measurement conditions

En_op          = 1200;                          % Energy at which the holographic imaging simulation is performed
lambda_op      = (1239.842/En_op)*nanometers;   % Operation lambda
f_op           = f*lambda/lambda_op;            % Operation focal length

%% OSA parameters

OsaD        = 25*microns;                % OSA diameter
OsaR        = OsaD/2;

OSA         = X.^2+Y.^2<=(OsaR)^2;
OSApos      = 8*f_op/9;                  % OSA position relative to f operation

%% Assign Material Properties to the Lens
% Material Parameters

t           = 280*nanometers;                                               % Thickness of the lens

lens_material   = 'Au';

n_lens      = RefractiveIndex(En_op,lens_material);                         % Pull the refractive index data 

delta_lens  = real(n_lens);                                                 % Coefficients for FZP material
beta_lens   = imag(n_lens);

illumination_intensity = 1;
binary_CGH  = support.*(illumination_intensity.*((binary_CGH<DC).*1+(binary_CGH>=DC).*exp(-1i*2*pi*t*delta_lens/lambda_op).*exp(-2*pi*t*beta_lens/lambda_op)));


%% Assign Material Properties to the Sample

experiment      = 1;                            % 1 - Structural Imaging
                                                % 2 - Magnetic Imaging
                                                               
switch experiment
    
    case 1
t_sample        = 180*nanometers;               % Thickness of the sample
sample_material = 'Au';
n_sample      = RefractiveIndex(En_op,sample_material);

delta_sample  = real(n_sample);                 % Coefficients for FZP material
beta_sample   = imag(n_sample);

Sam_size      = 20*microns;                     % Sample width
SamNx         = size(test,1);
Samdx         = Sam_size./SamNx;

test          = imresize(test,SamNx/size(test,1));
sample        = (test<0.5).*1+(test>=0.5).* exp(-1i*2*pi*t_sample*delta_sample/lambda_op).*exp(-2*pi*t_sample*beta_sample/lambda_op);


    case 2
% CoFeB Magnetic scattering factors are taken from Daan Boltje's MSc
% thesis: Voltage Induced Near Interface Changes of the Magnetocrystalline Anisotropy Energy: 
% A study by X-ray Resonant Techniques in Combination with Conventional Magnetometry (Universität Stuttgart, 2017).        


Sam_size      = 10*microns;                         % Sample width
SamNx         = size(test,1);
Samdx         = Sam_size./SamNx;

test          = imresize(test,SamNx/size(test,1));
test(test==0) = -1;
mz            = test;        
        
t_sample        = 50*nanometers;                    % Sample thickness


density         = 8.9;                              % g/cm^3 according to CXRO Database

Na              = 6.02214129*1e23;                  % Avogadro's Number #/mol
CFe             = 0.055;                            % Molar concentration mol/cm^3 of Fe in CoFeB as reported in Daan 2017 Fig. 7.6

na              = Na*CFe;                           % number density of Fe atoms in #/cm^3
na              = na*1e6;                           % number density of Fe atoms in #/m^3

re              = 2.8179403267*1e-15*meters;

% Electronic Structure Factors for Fe in CoFeB

fc1           = -48.32;
fc2           = 23.12;

% Magnetic Structure Factors for Fe in CoFeB

fm1           = 7.467;
fm2           = -5.483;


delta_sample  = (1/2*pi)*(na*re*lambda_op^2).*(fc1+fm1.*mz);         
beta_sample   = (1/2*pi)*(na*re*lambda_op^2).*(fc2+fm2.*mz);


sample        = exp(-1i*2*pi*t_sample*delta_sample/lambda_op).*exp(-2*pi*t_sample*beta_sample/lambda_op);

        
end

%% Apply a Beam Stop

switch beam_stop
    
    case 1
BSD         = 5*microns;                               % BS diameter in um
BS          = X.^2+Y.^2>=(BSD/2)^2;
aperture    = X.^2+Y.^2<=R^2;
    case 2
BS          = 1;                                        % No BS
aperture    = 1;                                        % No aperture

end

binary_CGH  = binary_CGH.*BS.*aperture.*support;


%% Boundary Conditions

S           = L;                        % Simulation size at critical condition that satisf. dx=lambda*f/S
ds          = dx;
s           = -S/2:ds:S/2-ds;           % New simulation coordinates

[Sx,Sy]     = meshgrid(s);
sim_size    = length(s);
Sim         = zeros(sim_size);
ScaledOSA   = zeros(sim_size);


margin1     = (length(s)-length(x))/2;  % Margins for the lens



G=zeros(Nx);
Sim((1+margin1):(length(s)-margin1),(1+margin1):(length(s)-margin1))=binary_CGH;

margin3     = margin1;                  % Margins for OSA, same as the lens

ScaledOSA((1+margin3):(length(s)-margin3),(1+margin3):(length(s)-margin3))=OSA;

binary_CGH  = Sim;

OSA         = ScaledOSA;


%% Propagate the CGH to Focal Plane with the option of an OSA in between

switch Use_OSA
%% Insert OSA
    case 1
% Propagate to the OSA plane 
propOSA      = propRS2(binary_CGH,S,lambda_op,OSApos);

% filter spurious orders
intermed    = propOSA.*OSA;

% Propagate to the focal plane
P           = propRS2(intermed,S,lambda_op,f_op-OSApos);

%% No OSA
    case 2 
        
P           = propRS2(binary_CGH,S,lambda_op,f_op);

end

%% Calculate the Diffraction Efficiency and the Figelity at the fcal plane

Focalmask   = circle(15*Ref_spot_D,0,0,x)+Fov_spot;
Focus_int   = Focalmask.*abs(P).^2;

DiffEff     = sum(sum(Focus_int))/sum(sum(support.^2))*100      % Diffraction efficiency

Fidelity    = sum(sum(Focus_int))/sum(sum(abs(P).^2))*100       % Ratio of intensity in the desired target pattern to the whole focal plane pattern.

return
%% Save 3D propagation field

% This section is to plot the intensity profile as a function of propagation
% distance. One can skip this whole section and simulate the second half of the experiment.
%

% Calculate the beam caustics

counter2 = 1;                       % Initialize the counter

dz       = 50*microns;              % 25 microns give good results
for z=0:dz:1.50*f_op                % Loop through all z-sampling points
    field=propRS2(binary_CGH,S,lambda_op,z);
    m(:,counter2)= diag(fliplr(field));
    counter2=counter2+1
end
z           = 0:dz:1.5*f_op;        % Create a z-vector for plotting
yz_slice    = m;

yz_diag     = sqrt(2)*S;
dx_diag     = sqrt(2)*dx;
x_diag      = -yz_diag/2:dx_diag:yz_diag/2-dx_diag;


% Plot the beam caustics in the vicinity of the focal plane
% Phase plot
figure(1)
imagesc(z.*1e3,x_diag.*1e6,angle(yz_slice));
title('$\angle{diag(E)}$','Interpreter','latex')
colormap(cmap)
xlabel 'Propagation Distance [mm]'
ylabel 'Aperture Distance [\mum]'
pbaspect([16 9 1])
xlim([11 13])
ylim([-7. 7.5])

% Amplitude Plot
figure(2)
imagesc(z.*1e3,x_diag.*1e6,sqrt(abs(yz_slice)));
title('$\sqrt{diag(E)}$','Interpreter','latex')
colormap jet
xlabel 'Propagation Distance [mm]'
ylabel 'Aperture Distance [\mum]'
pbaspect([16 9 1])
xlim([11 13])
ylim([-7 7.5])

% Complex Plot
figure(3)
dcolor(z.*1e3,x_diag.*1e6,flipud(yz_slice),'cfun',9);
title('Complex $diag(E)$','Interpreter','latex')
% xlim([10.5 13.5])
% ylim([-7 7])
xlabel 'Propagation Distance [mm]'
ylabel 'Aperture Distance [\mum]'
pbaspect([16 9 1])
xlim([11 13])
ylim([-7. 7.5])

return

%% Plot the optics and sample beam interatcions

figure(4)
subplot(1,4,1)
imagesc(s*1e6,s*1e6,abs(binary_CGH))
subplot(1,4,2)
imagesc(s*1e6,s*1e6,angle(P))
subplot(1,4,3)
imagesc(s*1e6,s*1e6,log10(abs(P)))
subplot(1,4,4)
dcolor(s*1e6,s*1e6,flipud(P),'cfun',2)

% Clear cache because these are not used below and they occupy a huge
% memory space

clear ScaledOSA CGH CGH2 hologram A B C D G intermed lens propOSA source target x y X Y support OSA Ref_spot Recons mask1 mask2 target_abs target_phase bb holo holo2 C1 C2 Fov_spot

%% Simulate the CGXH holography experiment;
% Sections below this point simulates the holography experiment downstream the sample 

% Upscale the P at f to match the SEM image sampling (because SEM images 
% and simulated optics don't have the same dimensions)

ds          = Samdx;
s           = -S/2:ds:S/2-ds; % New simulation coordinates
[Sx,Sy]     = meshgrid(s);
Newsim_size = length(s);

P           = imresize(P,Newsim_size/size(P,1));

margin2     = (length(s)-size(test,1))/2; % Margins for the sample (sample size is 20 um for Siemens star with 1000 pix)
margin3     = (length(s)-size(test,2))/2;

% Insert the synthetic sample into the new simulation support
% This part can be used to shift sample to image different parts
% Default values are 0 shift.

ScaledSample= zeros(Newsim_size);
shiftY      = 0*microns;
shiftXpix   = round(shiftY./dx);
shiftX      = 0*microns;
shiftYpix   = round(shiftX./dx);
ScaledSample((1+margin2+shiftXpix):(length(s)-margin2+shiftXpix),(1+margin3+shiftYpix):(length(s)-margin3+shiftYpix))=sample;

illumination_int = 1;                % Pre-multiplier for intensity per pixel


exit_wave   = ScaledSample.*P;       % This is the exit-wave function that creates the far-field interference pattern

% Plot results
figure(5)
subplot(1,2,1)
imagesc(s*1e6,s*1e6,(test))
xlabel 'Distance \mum'
ylabel 'Distance \mum'
title 'Siemens Star Test Object'
axis square
axis tight
colormap gray
subplot(1,2,2)
dcolor(s*10^6,s*10^6,flipud(exit_wave),'cfun',2)
xlabel 'Distance \mum'
ylabel 'Distance \mum'
title 'Complex Exit Wave'
axis square
axis tight
xlim([-6.5 6.5])
ylim([-6.5 6.5])

% Clear some more variables to reduce memory allocation
clear Sim ScaledSample

%% Simulate Fraounhofer diffraction pattern with different side lengths;
clear interference_2 off_axis_recon_noisy

det_z                   = 16*cm;                    % sample to detector distance

% Detector plane side length and sampling

det_L                   = (lambda_op*det_z)/ds;
dx_det                  = lambda_op*det_z/S;
det_coord_x             = lambda_op*det_z*(-1/(2*ds):1/S:1/(2*ds)-1/S);

% Detector parameters (taken from pnCCD of MAXYMUS at BESSY II)

detPix                  = 48*microns;
detPixNum               = 264;
detSize                 = detPixNum*detPix;

interference_2          = abs(propFF(exit_wave,S,lambda_op,det_z)).^2;

interference_3          = (interference_2./max(max(interference_2)));


% Calculate the detector plane interference pattern with or without noise

% Poisson Noise - Detector Shot Noise

max_photons_per_pix     = 1000;
interference_3          = max_photons_per_pix.*interference_3.*1e-12;

interference_4          = imnoise(interference_3,'poisson');


%% Fourier Transform Holography - Reconstruct the complex exit wave function

off_axis_recon_noisy    = fftshift(ifft2(fftshift(interference_4)));

% Plot the reconstruction results in amplituede and phase

figure(6)
imagesc(s*1e6,s*1e6,sqrt(abs(off_axis_recon_noisy)));
colormap gray
xlabel 'Distance \mum'
ylabel 'Distance \mum'
title('$\sqrt{|E|}$','Interpreter','latex')
axis square
axis tight
xlim([0.5 3.5])
ylim([-3.5 -0.5])
caxis([0 1e-7])

figure(7)
imagesc(s*1e6,s*1e6,angle(off_axis_recon_noisy))
colormap(cmap)
xlabel 'Distance \mum'
ylabel 'Distance \mum'
title('\angle{E}','Interpreter','latex')
axis square
axis tight
xlim([0.5 3.5])
ylim([-3.5 -0.5])
% caxis([-pi/2 pi/2])

% Calculate the Fresnel numbers

FresnelNumber1=CGXH_w.^2/(lambda_op/f_op)
FresnelNumber2=fov_FW^2/(lambda_op*det_z)


%% Calculate the color wheel and plot it

rr     = sqrt(Sx.^2+Sy.^2);
abs_wheel = exp(1i.*2.*pi.*rr);
wheel  = angle(abs_wheel).*exp(1i.*atan2(Sx,Sy));

cwheel = cat(3,angle(wheel),abs(wheel),abs(wheel));

figure(8)
dcolor(wheel,'cfun',2)

%%%%%%%%%%%%%%%%%%%%%%%%%%% End of Script %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%